package com.zang.blogz.utils;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

/**
 * @Author: 橘贝贝
 * @Description:  限流工具类
 * @DateTime: 2022/6/30 15:19
 **/
@Component
public class LimitUtil {

    /**
     * redis key前缀
     */
    private static final String SLIDING_WINDOW = "sliding_window_";

    @Autowired
    private RedisTemplate redisTemplate;


    public boolean canAccess(String key, int windowInSecond, long maxCount) {

        key = SLIDING_WINDOW + key;


        Long count = redisTemplate.opsForZSet().zCard(key);
        if (count == maxCount){
            //先清空s过时v再统计一遍
            long currentMs = System.currentTimeMillis();
            long windowStartMs = currentMs - windowInSecond * 1000;

            ZSetOperations zSetOperations = redisTemplate.opsForZSet();
            zSetOperations.removeRangeByScore(key, 0, windowStartMs);

            count = redisTemplate.opsForZSet().zCard(key);
        }
        if (count < maxCount) {

            increment(key, windowInSecond);
            return true;
        } else {

            return false;
        }
    }

    public void increment(String key, Integer windowInSecond) {

        long currentMs = System.currentTimeMillis();
        long windowStartMs = currentMs - windowInSecond * 1000;

        ZSetOperations zSetOperations = redisTemplate.opsForZSet();
        zSetOperations.removeRangeByScore(key, 0, windowStartMs);

        zSetOperations.add(key, String.valueOf(currentMs), currentMs);

        redisTemplate.expire(key, windowInSecond, TimeUnit.SECONDS);
    }

}
